import pandas as pd
import os

import sqlite3 as sql
# 读取文件
def read_code(code, year):
    data = pd.read_csv('./stocks/%s/%s.csv' % (year, code), index_col=0)
    return data


def read_all(code):
    temp = pd.concat([read_code(code, str(i)) for i in range(2013, 2021)], ignore_index=True)
    return temp


def count_high(df):
    # 判断第七天是否是第一天的1.6倍
    if df.empty:
        return
    # 直接group
    df['七天翻倍'] = df.rolling(window=7)['close'].apply(lambda x: x[6] / x[0], raw=True)
    revert_close = df['close'].sort_index(ascending=False)
    for i in range(1, 11):
        df['%d天后' % i] = revert_close\
            .rolling(window=(i+1))\
            .apply(lambda x: x[0] / x[i], raw=True)\
            .sort_index(ascending=True)
    return df[lambda x: x['七天翻倍'] > 1.6][lambda x:x['1天后'] < 1]

def _map():
    file_path = 'stocks/2018'
    datas = []
    for top, d, files in os.walk(file_path):
        for file in files:
            stock = file[:-4]
            data = count_high(read_all(stock))
            print('complete %s' % stock)
            if data is None or data.empty:
                continue
            datas.append(data)
    tot = pd.concat(datas)
    print(tot)
    tot.to_excel('consume/consume2013-continue-new.xlsx', index=False)


def reduce():
    # 合并所有map

    file_path = 'alys'
    for top, d, files in os.walk(file_path):
        if len(d) != 0:
            continue
        i = 0
        total = len(files)
        datas = []
        year = top[-4:]

        for file in files:
            print('开始 %s %s' % (top, file))
            stock = file[:-4]
            year = top[-4:]
            i += 1
            data = pd.read_csv('%s/%s' % (top, file), index_col='start_date')[lambda x: x.continue_day >= 5]
            if data.empty or len(data) == 0:
                continue
            data['code'] = stock
            datas.append(data)
            print('完成了 %s, %d/%d' % (top, i, total))
        pd.concat(datas).to_csv('consume/aly_%s.csv' % year)


if __name__ == '__main__':
    _map()
    # sql.connect('a.db')